import torch
import pandas as pd
import os
from model.attnSleepModel import AttnSleep

def main():
    # 加载模型 - 使用trained_model.pth
    model_path = os.path.join(os.path.dirname(__file__), 'trained_model.pth')
    model = AttnSleep()
    model.load_state_dict(torch.load(model_path))
    model.eval()

    # 读取CSV文件的第一行数据（包含表头）
    csv_path = os.path.join(os.path.dirname(__file__), 'segment_0.csv')
    data = pd.read_csv(csv_path, header=0, nrows=1)  # 修改header=0表示第一行是表头

    # 将数据转换为适合模型输入的格式
    input_data = data.to_numpy().reshape(1, 1, -1)
    input_tensor = torch.tensor(input_data, dtype=torch.float32)

    # 进行预测
    with torch.no_grad():
        outputs = model(input_tensor)
        _, predicted_class = torch.max(outputs, 1)

    # 输出预测结果
    print(f'Predicted class: {predicted_class.item()}')

if __name__ == "__main__":
    main()
